CIKM 2021 | 基于池化结构搜索的图分类

论文标题：

Pooling Architecture Search for Graph Classification

论文地址：

https://arxiv.org/abs/2108.10587

代码地址：

https://github.com/AutoML-Research/PAS

备注：该论文已经被数据挖掘会议 CIKM 2021 接收，欢迎大家关注。如有任何问题，欢迎联系 [e-mail]。

简介

近年来 GNN（Graph Neural Network）成为工业界和学术界非常热门的一个研究方向。在图分类任务中，不同的池化操作应用在了 graph-level 的特征学习中。但是面对多样的数据集和任务，没有任何一个方法能够稳定取得 SOTA 效果。最近，斯坦福大学 Jure 教授团队在 NeurIPS 2020 的工作上也指出了这一点 [1]。

基于此，我们结合自动化机器学习（AutoML）[2-5] 来获取任务自适应的池化图神经网络。具体来说，我们提出了一个基于图分类任务的统一框架，并在此基础上设计了新的搜索空间（search space）。为效率起见，我们提出了一个 coarsening strategy 来解决池化操作的松弛问题，从而可微搜索算法（differentiable search algorithm）可以被应用到池化结构搜索中。我们的池化结构搜索方法（Pooling Architecture Search，简记为 PAS）能够在取得 SOTA 表现的同时还十分高效。

背景介绍

图分类任务有大量的应用场景，比如蛋白质/化学分子性质预测 [6]，社交网络分析 [7]，文档分类 [8] 等，如图 1 所示。

▲ 图1：图分类的应用场景示例

相比较于 node-level 的任务，图分类任务需要学习 graph-level 的特征表示，池化操作也由此被提出。一个简单直接的方案就是取所有节点的的特征均值或者和，这类操作被称为全局池化操作（global pooling）；这类操作学到的是单层次的 graph-level 特征，无法捕获图结构中的层次化信息。层次化的池化（hierarchical pooling）方案就是为了解决这个问题，它通过生成越来越粗粒度的图来提取层次化的 graph-level 信息。

现有的图分类方案提供了不同的 global 和 hierarchical 的池化方案，比如 DGCNN [10]，DiffPool [11]，SAGPool [12]，ASAP [13]，Graph U-Net [14] 等。现有的 GNN 结合神经网络结构搜索（neural architecture search，简记为NAS）的方案，都是专注于 node-level 的工作，比如 GraphNAS [15]， SNAG [16]，SANE [17] 等，缺少 graph-level 的探索，方案 [18] 考虑了 global pooling 的设计，但是它基于演化算法（Evolutionary algorithm，简记为 EA），耗时长效率低。这些方案缺少了对池化结构的重要部分的探索，即缺少对 hierarchical pooling 的设计。基于上述，如何能快速高效的设计数据自适应的图分类任务 GNN 是一个大的挑战。

本次工作的方法

为了解决上述的问题，我们提出了新的解决方案 PAS（Pooling architecture search），其中包括一个统一框架，基于该框架的设计的搜索空间，以及考虑到效率而设计的 coarsening strategy。

从现有的图分类 GNN 中，我们提取了四个重要的模块，聚合模块（Aggregation Module），池化模块（Pooling Module），读出模块（Readout Module）和融合模块（Merge Module）。以图 1（b）的两层结构为例，针对输入的图 ，聚合模块更新了节点自身信息，结果可以表示为 ，随后池化模块产生粗粒度的图 。

这两个操作的示意图展示在（a）中，第二层的聚合模块和池化模块也是同样的操作。三个读出模块用来读出每一层的 graph-level 特征 ，最后聚合模块基于生成的特征 生成最终的 graph-level 特征 。

▲ 图2：PAS方案示意图

基于这个统一框架，我们提出了一个 NAS 方案来搜索池化结构。按照 NAS 的一贯设计，我们给出了搜索空间中各个模块的操作集合，如表 1 所示。

▲ 表1：搜索空间

在本文中，池化操作的过程可以被统一为如下公式。其中 存储着节点的得分，由打分函数 得到，之后 函数选择出分数排名前 的节点。粗粒度图的邻接矩阵 和节点特征矩阵 可以按照对应公式采样得到。基于这个公式，我们提供了 10 种池化操作，详情请见论文。

可微搜索算法中，一项重要的设计就是松弛函数。如下所示，离散空间可以被松弛为连续空间，其中  是集合 中第 个操作 的权重。

基于这个松弛函数，聚合模块、读出模块和融合模块的结果可以按照如下公式计算。

但是这对池化操作来说是不可实现的，如下图所示，池化操作 和 产生了不同的节点集合 和 ，此时粗粒度图中的特征矩阵都是 ，虽然具有相同的矩阵大小，但是因为节点集合不匹配所以无法直接加权求和。

我们设计了新的 coarsening strategy 来解决这个问题，在产生了候选的 个节点后，那些没有选中部分记为 0。也就是说，下图中未被选中的节点（灰色节点）特征为 0，未被选中的边（灰色的边）权重为 0。当前策略下生成的粗粒度图，具有和输入图一样的“形状”，因此池化模块的结果可以按照下述公式计算。

▲ 图3：池化模块

在设计的 coarsening strategy 基础上，整个框架可以用梯度下降来优化。搜索完成后，取出每个模块中权重最大的操作，基于此我们可以得到搜索的池化结构。

实验

该工作在六个常用的图分类数据集上做了验证。

4.1 PAS的实验效果

▲ 表2：实验效果对比

▲ 图4：不同数据集上搜索的GNN结构

▲ 图5：准确率和模型参数的对比

如表 2 所示，PAS 比现有的两类池化方案和 NAS 方案的效果更好，这也体现了我们提出的方案的有效性。搜索出的 GNN 结构如图 4 所示，不同数据集上的结构各不相同，并且在 COX2 数据集上，搜索的 GNN 退化为一个 global pooling 的方案。图 5 中展示了测试集上准确率和模型参数的对比，PAS 搜索出来的结构能够用更少的参数获得更高的结果。综合这些实验结果，PAS 的效率和效果是显而易见的。

4.2 消融实验（Ablation Study）

针对 PAS 的四个模块，我们设计了对应的消融实验来验证四个模块的重要性。如表 3 所示，这些方案的效果低于 PA 本身，这展示了我们设计的统一框架的有效性。

▲ 表3：消融实验中的效果对比

4.3 效率对比

如下图所示，我们对比了不同方案的搜索空间大小和搜索耗时。圆圈大小代表了搜索空间的大小，不同颜色代表了不同方案。PAS 的搜索空间大小介于 GraphNAS 和 SNAG 之间，但是具有最小的搜索耗时和最高的效果。结合图 5，PAS 的高效性显而易见。

▲ 图6：搜索耗时和搜索空间的对比

结论

本文提出了一种自动设计池化结构的方案 PAS。PAS 中包含一个新的统一框架，它由四个常用的模块组成。在此基础上，我们提供了对应的搜索空间并设计了 coarsening strategy 来松弛这个搜索空间。基于此，可微搜索这种高效的搜索算法能被应用到池化结构搜索中。我们设计了充分的实验来验证搜索空间和搜索算法这两部分，结果显示 PAS 比现有的方案结果更好效率更高。

本组其他相关工作

本文探索了 NAS 在图分类任务中的应用，除此之外，本组还有以下新工作：

• AutoSF: Searching scoring functions for knowledge graph embedding. ICDE 2020. (AutoSF)

• Interstellar: Searching Recurrent Architecture for Knowledge Graph Embedding. NeurIPS 2020. (Interstellar)

• Search to aggregate neighborhood for graph neural network. ICDE 2021. (SANE)

• Simplifying Architecture Search for Graph Neural Network. CIKM-CSSA 2020 (SNAG)

• Searching to Sparsify Tensor Decomposition for N-ary relational data. WWW 2021. (S2S)

• DiffMG: Differentiable Meta Graph Search for Heterogeneous Graph Neural Networks. KDD 2021. (DiffMG)

• TabGNN: Multiplex Graph Neural Network for Tabular Data Prediction. KDD-DLP 2021. (TabGNN)

未来工作

在未来的工作中，我们打算结合本组的工作 SANE 和 SNAG，加强 PAS 中对聚合操作的探索，并打算深入探索搜索出来的 GNN 和图性质之间的关系，同时尝试将 PAS 应用在更大的图数据集中，比如 OGB。

参考文献

[1] Design Space for Graph Neural Networks. NeurIPS 2020.

[2] Neural architecture search with reinforcement learning. ICLR 2017

[3] Regularized evolution for image classifier architecture search. AAAI 2019

[4] DARTS: Differentiable architecture search. ICLR 2019

[5] SNAS: stochastic neural architecture search. ICLR 2019

[6] Distinguishing enzyme structures from non-enzymes without alignments. JMB 2003

[7] How powerful are graph neural networks? ICLR 2019

[8] Text categorization as a graph classification problem. ACL 2015

[9] 图1左：https://www.mdpi.com/2078-2489/1/2/60/htm, 中:https://medium.com/analytics-vidhya/social-network-analytics-f082f4e21b16, 右：论文[8]。

[10] An end-to-end deep learning architecture for graph classification. AAAI 2018.

[11] Hierarchical graph representation learning with differentiable pooling. NeurIPS 2018.

[12] Self-Attention Graph Pooling. ICML 2019

[13] ASAP: Adaptive Structure Aware Pooling for Learning Hierarchical Graph Representations. AAAI 2020

[14] Graph u-nets. ICML 2019.

[15] Graphnas: Graph neural architecture search with reinforcement learning. IJCAI 2020.

[16] Simplifying Architecture Search for Graph Neural Network. CIKM-CSSA 2020

[17] Search to aggregate neighborhood for graph neural network. ICDE 2021

[18] Graph Neural Network Architecture Search for Molecular Property Prediction. Arxiv 2020

特别鸣谢

感谢 TCCI 天桥脑科学研究院对于 PaperWeekly 的支持。TCCI 关注大脑探知、大脑功能和大脑健康。

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：[e-mail] 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·